"""
中文垃圾邮件检测
"""
import pandas as pd 
import jieba.posseg as pseg 
from sklearn.feature_extraction.text import TfidfVectorizer     #  TF-IDF向量化
from sklearn.naive_bayes import MultinomialNB                   #  朴素贝叶斯分类器 
from sklearn.svm import SVC                                     #  支持向量机分类器  
from sklearn.ensemble import RandomForestClassifier             #  随机森林分类器  
from sklearn.model_selection import cross_val_score       # 数据集划分
import re                       # 正则
import pickle                   # 模型持久化
from config import * 



class SpamDetector:
    def __init__(self):
        self.vectorizer = None; 
        self.model = None;
        self.stop_words = self._load_stop_words() 

    def _load_stop_words(self):
        # 加载停用词 
        return CHINESE_STOP_WORDS
    
    def preprocess_text(self, text):
        # 文本预处理 
        text = re.sub(r'[^\u4e00-\u9fa5]', '', text)  # 只保留中文字符
        words = pseg.cut(text)  # 分词并词性标注

        meaningful_words = []           # 存储有意义的词
        for word, flag in words:
            if (len(word) >= 2 and 
                word not in self.stop_words and 
                flag in ['n', 'v', 'a', 'nr', 'ns', 'nt', 'nz', 'vn', 'an']):
                meaningful_words.append(word)

        return ' '.join(meaningful_words)
    
    def create_extended_dataset(self):
        # 创建扩展数据集
        SPAM_EMAILS = spam_emails
        NORMAL_EMAILS = normal_emails

        emails = SPAM_EMAILS + NORMAL_EMAILS 
        labels = ['spam'] * len(spam_emails) + ['ham'] * len(normal_emails)

        return pd.DataFrame({'text': emails, 'label': labels})   
    

    def train_with_multiple_models(self):
        # 使用多个模型进行训练
        # 创建数据集
        data = self.create_extended_dataset()

        # 预处理文本
        data['text'] = data['text'].apply(self.preprocess_text)  

        # 特征提取
        self.vectorizer = TfidfVectorizer(**MODEL_CONFIG)
        X = self.vectorizer.fit_transform(data['text'])
        y = data['label'].map({'ham': 0, 'spam': 1})

        # 比较多个模型 
        models = {
            'Naive Bayes': MultinomialNB(),
            'SVM': SVC(kernel='linear', probability=True),
            'Random Forest': RandomForestClassifier(n_estimators=100)
        }
        best_score = 0
        best_model_name = ""

        for name, model in models.items():  
            scores = cross_val_score(model, X, y, cv=5)
            avg_score = scores.mean()
            print(f"{name} 平均准确率: {avg_score:.4f}")

            if avg_score > best_score:
                best_score = avg_score
                best_model_name = name
                self.model = model

        print(f"最佳模型: {best_model_name}，平均准确率: {best_score:.4f}")

        # 训练最佳模型 
        self.model.fit(X, y) 

        return data 
    

    def predict(self, text):
        # 预测垃圾邮件
        processed = self.preprocess_text(text)                  # 预处理
        vector = self.vectorizer.transform([processed])          # 向量化
        prediction = self.model.predict(vector)[0]              # 预测
        probability = self.model.predict_proba(vector)[0]       # 概率

        return {
            'text': text,
            'processed': processed,
            'prediction': '垃圾邮件' if prediction == 1 else '正常邮件',
            'spam_probability': probability[1],
            'confidence': max(probability)
        }
    

    def save_model(self, filepath='spam_model.pkl'):
        """保存模型"""
        model_data = {
            'vectorizer': self.vectorizer,
            'model': self.model,
            'stop_words': self.stop_words
        }
        with open(filepath, 'wb') as f:
            pickle.dump(model_data, f)
        print(f"模型已保存到: {filepath}")
    
    def load_model(self, filepath='spam_model.pkl'):
        """加载模型"""
        with open(filepath, 'rb') as f:
            model_data = pickle.load(f)
        
        self.vectorizer = model_data['vectorizer']

        print(self.vectorizer)  
        self.model = model_data['model']
        self.stop_words = model_data['stop_words']
        print("模型加载成功！") 

def main():
    detector = SpamDetector()
    
    # 训练模型
    print("开始训练模型...\n")
    detector.train_with_multiple_models() 
    
    detector.save_model()
    detector.load_model()

    # 测试邮件
    test_emails = [
        "恭喜您中奖一千万！请立即领取！",
        "明天开会请准时参加。",
        "免费送手机！点击领取！",
        "您的快递已到达，请取件。",
        "投资理财！月息30%！无风险！",
        "会议纪要请查收附件。"
    ]
    
    print("\n=== 测试结果 ===")
    for email in test_emails:
        result = detector.predict(email)
        print(f"\n邮件: {result['text']}")
        print(f"预测: {result['prediction']}")
        print(f"垃圾邮件概率: {result['spam_probability']:.4f}")
        print(f"置信度: {result['confidence']:.4f}")

if __name__ == "__main__":
    main()